package com.atguigu.utils;

import org.apache.flink.streaming.api.functions.sink.{TwoPhaseCommitSinkFunction, SinkFunction}
import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaProducer
import org.apache.flink.api.common.serialization.SimpleStringSchema
import org.apache.flink.api.common.functions.MapFunction
import java.util.Properties

class KafkaTransactionalSink(brokers: String, topic: String)
        extends TwoPhaseCommitSinkFunction[String, FlinkKafkaProducer[String]] {

        private var producer: FlinkKafkaProducer[String] = _

        override def beginTransaction(): FlinkKafkaProducer[String] = {
        // create a new instance of FlinkKafkaProducer with transactional.id set
        val kafkaProps = new Properties()
        kafkaProps.setProperty("bootstrap.servers", brokers)
        kafkaProps.setProperty("transactional.id", "my-transactional-id")
        val kafkaProducer = new FlinkKafkaProducer[String](topic, new SimpleStringSchema(), kafkaProps)
        kafkaProducer.initTransactions()
        producer = kafkaProducer
        producer.beginTransaction()
        producer
        }

        override def preCommit(transaction: FlinkKafkaProducer[String]): Unit = {
        // no-op
        }

        override def commit(transaction: FlinkKafkaProducer[String]): Unit = {
        producer.commitTransaction()
        }

        override def abort(transaction: FlinkKafkaProducer[String]): Unit = {
        producer.abortTransaction()
        }

        override def invoke(record: String, context: TwoPhaseCommitSinkFunction.Context): Unit = {
        // send the record to the producer
        producer.send(record)
        }
        }

object KafkaTransactionalSinkExample {
        def main(args: Array[String]): Unit = {
        val env = StreamExecutionEnvironment.getExecutionEnvironment
        val stream = env.fromElements("foo", "bar")

                // 使用自定义的kafka sink
        val kafkaSink = new KafkaTransactionalSink("localhost:9092", "my-topic")

        stream.map(new MapFunction[String, String] {
        override def map(value: String): String = value.toUpperCase()
        }).addSink(kafkaSink)

        env.execute("Kafka Transactional Sink Example")
        }
        }
